#include <cassert>
#include <cstring>
#include <iostream>
#include <unordered_map>
#include "visit.hpp"
#include "koopa.h"

// 存储类型为 koopa_raw_value 的有返回值的语句的位置
static std::unordered_map<koopa_raw_value_t, int> valueLocation;
// 栈空间总长度
static int totalStackSize = 0;
// 已使用栈空间长度
static int usedStackSize = 0;
// 当前函数是否保存了返回地址
static int isReturnAddressSaved = 0;
// 是否需要保存返回地址
static bool needToSaveReturnAddress = false;

// 处理 raw program
void Visit(const koopa_raw_program_t &program) {
    Visit(program.values);
    Visit(program.funcs);
}

// 处理 raw slice
void Visit(const koopa_raw_slice_t &slice) {
    for (size_t i = 0; i < slice.len; ++i) {
        auto element = slice.buffer[i];
        switch (slice.kind) {
            case KOOPA_RSIK_FUNCTION:
                Visit(reinterpret_cast<koopa_raw_function_t>(element));
                break;
            case KOOPA_RSIK_BASIC_BLOCK:
                Visit(reinterpret_cast<koopa_raw_basic_block_t>(element));
                break;
            case KOOPA_RSIK_VALUE:
                Visit(reinterpret_cast<koopa_raw_value_t>(element));
                break;
            default:
                assert(false);  // 不支持的类型
        }
    }
}

// 将值加载到寄存器
static void loadToRegister(const koopa_raw_value_t &value, const std::string &reg) {
    if (value->kind.tag == KOOPA_RVT_INTEGER) {
        std::cout << "  li " << reg << ", " << value->kind.data.integer.value << "\n";
    } 
    else if (value->kind.tag == KOOPA_RVT_FUNC_ARG_REF) {
        // 处理函数参数
        const auto& index = value->kind.data.func_arg_ref.index;
        if (index < 8) {
            std::cout << "  mv " << reg << ", a" << index << std::endl;
        } else {
            std::cout << "  li t6, " << totalStackSize + (index - 8) * 4 << std::endl;
            std::cout << "  add t6, t6, sp" << std::endl;
            std::cout << "  lw " << reg << ", 0(t6)" << std::endl;
        }
    }
    else if (value->kind.tag == KOOPA_RVT_GLOBAL_ALLOC) {
        // 处理全局变量
        std::cout << "  la t6, " << value->name + 1 << std::endl;
        std::cout << "  lw " << reg << ", 0(t6)" << std::endl;
    } else {
        std::cout << "  li t2, " << valueLocation[value] << "\n";
        std::cout << "  add t2, sp, t2\n";
        std::cout << "  lw " << reg << ", 0(t2)\n";
    }
}

// 计算栈空间大小并进行对齐
static void calculateStackSpace(const koopa_raw_function_t &func) {
    int localVarCount = 0;    // 局部变量数量
    int needsReturnAddress = 0; // 是否需要返回地址空间
    int argumentSpace = 0;      // 额外的参数空间

    for (size_t i = 0; i < func->bbs.len; ++i) {
        const auto& instructions = reinterpret_cast<koopa_raw_basic_block_t>(func->bbs.buffer[i])->insts;
        localVarCount += instructions.len;

        for (size_t j = 0; j < instructions.len; ++j) {
            auto instruction = reinterpret_cast<koopa_raw_value_t>(instructions.buffer[j]);

            // 处理不占用栈空间的指令类型
            if (instruction->ty->tag == KOOPA_RTT_UNIT) {
                localVarCount--;  // 减少单位指令的数量
            }

            // 处理调用指令
            if (instruction->kind.tag == KOOPA_RVT_CALL) {
                needsReturnAddress = 1;  // 需要为返回地址分配空间
                needToSaveReturnAddress = true;  // 需要保存 ra
                argumentSpace = std::max(argumentSpace, int(instruction->kind.data.call.args.len) - 8); // 参数超过8个时需要额外空间
            }
        }
    }

    // 计算栈帧的总大小，每个变量占 4 字节，栈空间对齐为 16 字节
    totalStackSize = (localVarCount + needsReturnAddress + argumentSpace) * 4;
    totalStackSize = (totalStackSize + 16 - 1) & (~(16 - 1));  // 16字节对齐
    usedStackSize = argumentSpace * 4; // 已使用的栈空间
}

// 处理函数
void Visit(const koopa_raw_function_t &func) {
    if(func->bbs.len == 0) return;
    std::cout << "  .text\n  .globl " << func->name + 1 << "\n" << func->name + 1 << ":\n";
    
    // 清除栈帧信息
    totalStackSize = 0;
    usedStackSize = 0;

    // 计算栈帧的大小
    calculateStackSpace(func);

    // 调整栈指针
    if (totalStackSize != 0) {
        if (totalStackSize <= 2047) {
            std::cout << "  addi sp, sp, -" << totalStackSize << "\n";
        } else {
            std::cout << "  li t0, " << totalStackSize << "\n";
            std::cout << "  sub sp, sp, t0\n";
        }
    }

    if (needToSaveReturnAddress) {
        std::cout << "  li t0, " << totalStackSize - 4 << std::endl;
        std::cout << "  add t0, t0, sp" << std::endl;
        std::cout << "  sw ra, 0(t0)" << std::endl;
        isReturnAddressSaved = 1;
        needToSaveReturnAddress = false;
    } else {
        isReturnAddressSaved = 0;
    }

    Visit(func->bbs);
    std::cout << std::endl;
}

// 处理基本块
void Visit(const koopa_raw_basic_block_t &bb) {
    if (strncmp(bb->name + 1, "entry", 5))
        std::cout << bb->name + 1 << ":" << std::endl;
    Visit(bb->insts);
}

// 处理指令
void Visit(const koopa_raw_value_t &value) {
    switch (value->kind.tag) {
        case KOOPA_RVT_INTEGER:
            Visit(value->kind.data.integer);
            break;
        case KOOPA_RVT_ALLOC:
            valueLocation[value] = usedStackSize;
            usedStackSize += 4;
            break;
        case KOOPA_RVT_LOAD:
            Visit(value->kind.data.load, value);
            break;
        case KOOPA_RVT_STORE:
            Visit(value->kind.data.store);
            break;
        case KOOPA_RVT_BINARY:
            Visit(value->kind.data.binary, value);
            break;
        case KOOPA_RVT_BRANCH:
            Visit(value->kind.data.branch);
            break;
        case KOOPA_RVT_JUMP:
            Visit(value->kind.data.jump);
            break;
        case KOOPA_RVT_CALL:
            Visit(value->kind.data.call, value);
            break;
        case KOOPA_RVT_GLOBAL_ALLOC:
            Visit(value->kind.data.global_alloc, value);
            break;
        case KOOPA_RVT_RETURN:
            Visit(value->kind.data.ret);
            break;
        default:
            break;
    }
}

// 处理 integer 类型指令
void Visit(const koopa_raw_integer_t &integer) {
    std::cout << "  li a0, " << integer.value << "\n";
}

// 处理 load 指令
void Visit(const koopa_raw_load_t &load, const koopa_raw_value_t &value) {
    loadToRegister(load.src, "t0");
    if (value->ty->tag != KOOPA_RTT_UNIT) {
        valueLocation[value] = usedStackSize;
        usedStackSize += 4;
        std::cout << "  li t1, " << valueLocation[value] << "\n";
        std::cout << "  add t1, sp, t1\n";
        std::cout << "  sw t0, 0(t1)\n";
    }
}

// 处理 store 指令
void Visit(const koopa_raw_store_t &store) {
    loadToRegister(store.value, "t0");
    std::cout << "  li t1, " << valueLocation[store.dest] << "\n";
    std::cout << "  add t1, sp, t1\n";
    std::cout << "  sw t0, 0(t1)\n";
}

// 处理 binary 指令
void Visit(const koopa_raw_binary_t &binary, const koopa_raw_value_t &value) {
    loadToRegister(binary.lhs, "t0");
    loadToRegister(binary.rhs, "t1");

    switch (binary.op) {
        case KOOPA_RBO_NOT_EQ:
            std::cout << "  xor t0, t0, t1\n  snez t0, t0\n";
            break;
        case KOOPA_RBO_EQ:
            std::cout << "  xor t0, t0, t1\n  seqz t0, t0\n";
            break;
        case KOOPA_RBO_GT:
            std::cout << "  sgt t0, t0, t1\n";
            break;
        case KOOPA_RBO_LT:
            std::cout << "  slt t0, t0, t1\n";
            break;
        case KOOPA_RBO_GE:
            std::cout << "  slt t0, t0, t1\n  xori t0, t0, 1\n";
            break;
        case KOOPA_RBO_LE:
            std::cout << "  sgt t0, t0, t1\n  xori t0, t0, 1\n";
            break;
        case KOOPA_RBO_ADD:
            std::cout << "  add t0, t0, t1\n";
            break;
        case KOOPA_RBO_SUB:
            std::cout << "  sub t0, t0, t1\n";
            break;
        case KOOPA_RBO_MUL:
            std::cout << "  mul t0, t0, t1\n";
            break;
        case KOOPA_RBO_DIV:
            std::cout << "  div t0, t0, t1\n";
            break;
        case KOOPA_RBO_MOD:
            std::cout << "  rem t0, t0, t1\n";
            break;
        case KOOPA_RBO_AND:
            std::cout << "  and t0, t0, t1\n";
            break;
        case KOOPA_RBO_OR:
            std::cout << "  or t0, t0, t1\n";
            break;
        case KOOPA_RBO_XOR:
            std::cout << "  xor t0, t0, t1\n";
            break;
        case KOOPA_RBO_SHL:
            std::cout << "  sll t0, t0, t1\n";
            break;
        case KOOPA_RBO_SHR:
            std::cout << "  srl t0, t0, t1\n";
            break;
        case KOOPA_RBO_SAR:
            std::cout << "  sra t0, t0, t1\n";
            break;
        default:
            break;
    }

    if (value->ty->tag != KOOPA_RTT_UNIT) {
        valueLocation[value] = usedStackSize;
        usedStackSize += 4;
        std::cout << "  li t1, " << valueLocation[value] << "\n";
        std::cout << "  add t1, sp, t1\n";
        std::cout << "  sw t0, 0(t1)\n";
    }
}

// 处理 branch 指令
void Visit(const koopa_raw_branch_t &branch) {
    loadToRegister(branch.cond, "t0");
    std::cout << "  bnez t0, DOUBLE_JUMP_" << branch.true_bb->name + 1 << std::endl;
    std::cout << "  j " << branch.false_bb->name + 1 << std::endl;
    std::cout << "DOUBLE_JUMP_" << branch.true_bb->name + 1 << ":" << std::endl;
    std::cout << "  j " << branch.true_bb->name + 1 << std::endl;
}

// 处理 jump 指令
void Visit(const koopa_raw_jump_t &jump) {
    std::cout << "  j " << jump.target->name + 1 << std::endl;
}

// 处理 call 指令
void Visit(const koopa_raw_call_t &call, const koopa_raw_value_t &value) {
    for (size_t i = 0; i < call.args.len; ++i) {
        auto arg = reinterpret_cast<koopa_raw_value_t>(call.args.buffer[i]);
        if (i < 8) {
            loadToRegister(arg, "a" + std::to_string(i));
        } else {
            loadToRegister(arg, "t0");
            std::cout << "  li t6, " << (i - 8) * 4 << std::endl;
            std::cout << "  add t6, t6, sp" << std::endl;
            std::cout << "  sw t0, 0(t6)" << std::endl;
        }
    }
    std::cout << "  call " << call.callee->name + 1 << std::endl;
    if (value->ty->tag != KOOPA_RTT_UNIT) {
        valueLocation[value] = usedStackSize;
        usedStackSize += 4;
        std::cout << "  li t1, " << valueLocation[value] << "\n";
        std::cout << "  add t1, sp, t1\n";
        std::cout << "  sw a0, 0(t1)\n";
    }
}

// 访问 global alloc 指令
void Visit(const koopa_raw_global_alloc_t &global_alloc, const koopa_raw_value_t &value) {
    if (!value->name) return;  // 检查 name 是否为空

    std::cout << "  .data\n";
    std::cout << "  .globl " << value->name + 1 << "\n";
    std::cout << value->name + 1 << ":\n";

    switch (global_alloc.init->kind.tag) {
        case KOOPA_RVT_ZERO_INIT:
            std::cout << "  .zero 4\n";
            break;
        case KOOPA_RVT_INTEGER:
            std::cout << "  .word " << global_alloc.init->kind.data.integer.value << "\n";
            break;
        default:
            assert(false);  // 不支持的初始化类型
    }
    std::cout << "\n";
}



// 处理 return 指令
void Visit(const koopa_raw_return_t &ret) {
    if (ret.value != nullptr)
        loadToRegister(ret.value, "a0");

    if (isReturnAddressSaved) {
        std::cout << "  li t0, " << totalStackSize - 4 << std::endl;
        std::cout << "  add t0, t0, sp" << std::endl;
        std::cout << "  lw ra, 0(t0)" << std::endl;
    }

    if (totalStackSize != 0) {
        if (totalStackSize <= 2047) {
            std::cout << "  addi sp, sp, " << totalStackSize << "\n";  // 出栈
        } else {
            std::cout << "  li t0, " << totalStackSize << "\n";
            std::cout << "  add sp, sp, t0\n";  // 出栈
        }
    }

    std::cout << "  ret\n";
}
